from gensim.models import Word2Vec  # conda install gensim==4.2.0
from random import choice

ABS_PATH = '古诗词.txt'  # 语料路径
WINDOW = 16  # 滑窗大小
MIN_COUNT = 60  # 过滤低频字
VECTOR_SIZE = 125  # 词向量维度
TOPN = 14  # 生成诗词的开放度


class Model:
    def __init__(self):
        # 语料读取
        with open(ABS_PATH, encoding='utf-8') as f:
            ls_of_ls_of_c = [list(line.strip()) for line in f]
        # 模型训练
        self.model = Word2Vec(sentences=ls_of_ls_of_c, vector_size=VECTOR_SIZE, window=WINDOW, min_count=MIN_COUNT)
        self.chr_dict = self.model.wv.index_to_key  # 字典

    def poem_generator(self, title, form):
        """古诗词生成"""
        seq = list(title)
        for i in range(form[0]):
            for _ in range(form[1]):
                chrs = self.model.predict_output_word(seq[-WINDOW:], max(TOPN, len(seq) + 1))
                chrs = [t[0] for t in chrs if t[0] not in ['，', '。']]
                char = choice([c for c in chrs if c not in seq[len(title):]])
                seq.append(char)
            seq.append('，' if i % 2 == 0 else '。')
        # 返回标题+主体
        length = form[0] * (form[1] + 1)
        title = '《%s》' % ''.join(seq[:-length])
        poem = ''.join(seq[-length:])
        return title + '\n' + poem


def poetize():
    form = {'五言绝句': (4, 5), '七言绝句': (4, 7), '对联': (2, 9)}
    m = Model()
    while True:
        title = input('输入标题：').strip()
        poem = m.poem_generator(title, form['五言绝句'])
        print('\033[031m%s\033[0m' % poem)  # red
        poem = m.poem_generator(title, form['七言绝句'])
        print('\033[033m%s\033[0m' % poem)  # yellow
        poem = m.poem_generator(title, form['对联'])
        print('\033[036m%s\033[0m' % poem)  # purple
        print()


if __name__ == '__main__':
    poetize()
